/* Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
        limitations under the License.
==============================================================================*/

#include "checkpoint.h"

#include <fcntl.h>
#include <omp.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

#include <iostream>

#include "ckpt_data_handler/feat_admit_n_evict_ckpt/feat_admit_n_evict_ckpt.h"
#include "ckpt_data_handler/key_count_map_ckpt/key_count_map_ckpt.h"
#include "ckpt_data_handler/key_freq_map_ckpt/key_freq_map_ckpt.h"
#include "file_system/file_system_handler.h"
#include "utils/common.h"
#include "error/error.h"
#include "utils/time_cost.h"

using namespace std;
using namespace MxRec;

void Checkpoint::SaveModel(string savePath, CkptData& ckptData, RankInfo& mgmtRankInfo, const vector<EmbInfo>& embInfo)
{
    processPath = savePath;
    rankId = mgmtRankInfo.rankId;
    deviceId = mgmtRankInfo.deviceId;
    rankSize = mgmtRankInfo.rankSize;
    useDynamicExpansion = mgmtRankInfo.useDynamicExpansion;
    mgmtEmbInfo = embInfo;

    auto fileSystemHandler = make_unique<FileSystemHandler>();
    fileSystemPtr = fileSystemHandler->Create(savePath);

    LOG_INFO("Start host side saving data.");
    LOG_DEBUG("Start to create save data handler.");
    SetDataHandler(ckptData);
    LOG_DEBUG("Start save data process.");
    SaveProcess(ckptData);
    LOG_INFO("Finish host side saving data.");
}

void Checkpoint::LoadModel(string loadPath, CkptData& ckptData, RankInfo& mgmtRankInfo, const vector<EmbInfo>& embInfo,
                           const vector<CkptFeatureType>& featureTypes)
{
    processPath = loadPath;
    rankId = mgmtRankInfo.rankId;
    deviceId = mgmtRankInfo.deviceId;
    useDynamicExpansion = mgmtRankInfo.useDynamicExpansion;
    mgmtEmbInfo = embInfo;

    auto fileSystemHandler = make_unique<FileSystemHandler>();
    fileSystemPtr = fileSystemHandler->Create(loadPath);

    LOG_INFO("Start host side loading data.");
    LOG_DEBUG("Start to create load data handler.");
    SetDataHandler(featureTypes);
    LOG_DEBUG("Start load data process.");
    LoadProcess(ckptData);
    LOG_INFO("Finish host side loading data.");
}

void Checkpoint::SetDataHandler(CkptData& ckptData)
{
    dataHandlers.clear();
    if (!ckptData.keyCountMap.empty()) {
        dataHandlers.push_back(make_unique<KeyCountMapCkpt>());
    }
    if (!ckptData.table2Thresh.empty() && !ckptData.histRec.timestamps.empty() &&
        !ckptData.histRec.historyRecords.empty()) {
        dataHandlers.push_back(make_unique<FeatAdmitNEvictCkpt>());
    }
    if (!ckptData.ddrKeyFreqMaps.empty() && !ckptData.excludeDDRKeyFreqMaps.empty()) {
        dataHandlers.push_back(make_unique<KeyFreqMapCkpt>());
    }
}

void Checkpoint::SetDataHandler(const vector<CkptFeatureType>& featureTypes)
{
    auto featAdmitNEvictHandler = [this] {
        dataHandlers.push_back(make_unique<FeatAdmitNEvictCkpt>());
    };
    auto ddrKeyFreqMapHandler = [this] {
        dataHandlers.push_back(make_unique<KeyFreqMapCkpt>());
    };
    auto keyCountMapHandler = [this] {
        dataHandlers.push_back(make_unique<KeyCountMapCkpt>());
    };
    map<CkptFeatureType, function<void()>> setCkptMap{{CkptFeatureType::FEAT_ADMIT_N_EVICT, featAdmitNEvictHandler},
                                                      {CkptFeatureType::DDR_KEY_FREQ_MAP, ddrKeyFreqMapHandler},
                                                      {CkptFeatureType::KEY_COUNT_MAP, keyCountMapHandler}};

    for (const auto& featureType : featureTypes) {
        setCkptMap.at(featureType)();
    }
}

void Checkpoint::SaveProcess(CkptData& ckptData)
{
    for (const auto& dataHandler : dataHandlers) {
        dataHandler->SetProcessData(ckptData);
        vector<string> embNames{dataHandler->GetEmbNames()};
        vector<CkptDataType> saveDataTypes{dataHandler->GetDataTypes()};
        MakeUpperLayerSaveDir();
        MakeDataLayerSaveDir(embNames, saveDataTypes, dataHandler);
        SaveDataset(embNames, saveDataTypes, dataHandler);
    }
}

void Checkpoint::MakeUpperLayerSaveDir()
{
    innerDirPath = processPath;
    MakeSaveDir(innerDirPath);
}

void Checkpoint::MakeDataLayerSaveDir(const vector<string>& embNames, const vector<CkptDataType>& saveDataTypes,
                                      const unique_ptr<CkptDataHandler>& dataHandler)
{
    for (const auto& embName : embNames) {
        auto dataDir{innerDirPath + dirSeparator + embName};
        MakeSaveDir(dataDir);

        for (const auto& saveDataType : saveDataTypes) {
            auto dataDirName{dataHandler->GetDataDirName(saveDataType)};
            auto datasetPath{dataDir + dirSeparator + dataDirName};
            MakeSaveDir(datasetPath);
        }
    }
}

void Checkpoint::MakeSaveDir(const string& dirName) const
{
    CheckFileSystemPtr();
    fileSystemPtr->CreateDir(dirName);
}

Checkpoint::EmbSizeInfo Checkpoint::GetEmbeddingSize(const string& embName)
{
    EmbSizeInfo embSizeInfo;
    for (const auto& embInfo : mgmtEmbInfo) {
        if (embInfo.name == embName) {
            embSizeInfo.embSize = embInfo.embeddingSize;
            embSizeInfo.extEmbSize = embInfo.extEmbeddingSize;
            return embSizeInfo;
        }
    }
    return embSizeInfo;
}

bool Checkpoint::CheckEmbNames(const string& embName)
{
    for (const auto& embInfo : mgmtEmbInfo) {
        if (embInfo.name == embName && embInfo.isSave) {
            return true;
        }
    }
    return false;
}

void Checkpoint::SaveDataset(const vector<string>& embNames, const vector<CkptDataType>& saveDataTypes,
                             const unique_ptr<CkptDataHandler>& dataHandler)
{
    for (const auto& embName : embNames) {
        if (!CheckEmbNames(embName)) {
            continue;
        }
        auto dataDir{innerDirPath + dirSeparator + embName};
        for (const auto& saveDataType : saveDataTypes) {
            auto datasetPath{dataDir + dirSeparator + dataHandler->GetDataDirName(saveDataType)};
            auto datasetDir{datasetPath + dirSeparator + datasetName + to_string(rankId) + dataFileType};

            LOG_DEBUG("Start getting data from handler to: {}", datasetDir);
            auto transData{dataHandler->GetDataset(saveDataType, embName)};

            LOG_DEBUG("Start saving data to: {}", datasetDir);
            WriteStream(transData, datasetDir, transData.datasetSize, saveDataType);
        }
    }
}

void Checkpoint::WriteStream(CkptTransData& transData, const string& dataDir, size_t dataSize, CkptDataType dataType)
{
    CheckFileSystemPtr();

    ssize_t writeBytesNum;
    if (int32TransSet.find(dataType) != int32TransSet.end()) {
        writeBytesNum =
            fileSystemPtr->Write(dataDir, reinterpret_cast<const char*>(transData.int32Arr.data()), dataSize);
    } else if (int64TransSet.find(dataType) != int64TransSet.end()) {
        writeBytesNum =
            fileSystemPtr->Write(dataDir, reinterpret_cast<const char*>(transData.int64Arr.data()), dataSize);
    } else if (dataType == CkptDataType::ATTRIBUTE) {
        writeBytesNum =
            fileSystemPtr->Write(dataDir, reinterpret_cast<const char*>(transData.attribute.data()), dataSize);
    } else {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::LOGIC_ERROR, "Unknown CkptDataType.");
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    if (writeBytesNum == -1) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::IO_ERROR,
                           StringFormat("Error: Save data failed, data type: %s. "
                                        "An error occurred while writing file: %s.",
                                        CkptDataTypeName(dataType).c_str(), dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
    if (writeBytesNum != dataSize) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::IO_ERROR,
                           StringFormat("Error: Save data failed. data type: %s. "
                                        "Expected to write %d bytes, but actually write %d bytes to file %s.",
                                        CkptDataTypeName(dataType).c_str(), dataSize, writeBytesNum, dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
}

void Checkpoint::LoadProcess(CkptData& ckptData)
{
    for (const auto& dataHandler : dataHandlers) {
        vector<string> embNames{};
        vector<string> dirNames{dataHandler->GetDirNames()};
        vector<CkptDataType> saveDataTypes{dataHandler->GetDataTypes()};
        innerDirPath = processPath;
        if (find(dirNames.begin(), dirNames.end(), ssdSymbol) != dirNames.end()) {
            embNames = GetTableLayerLoadDir();
        } else {
            embNames = GetEmbedTableNames();
        }
        LoadDataset(embNames, saveDataTypes, dataHandler, ckptData);
        dataHandler->GetProcessData(ckptData);
    }
}

vector<string> Checkpoint::GetEmbedTableNames()
{
    vector<string> loadTableNames;
    for (const auto& embInfo : mgmtEmbInfo) {
        if (embInfo.isSave) {
            loadTableNames.push_back(embInfo.name);
        }
    }

    return loadTableNames;
}

vector<string> Checkpoint::GetTableLayerLoadDir()
{
    vector<string> loadTableDir;
    CheckFileSystemPtr();
    loadTableDir = fileSystemPtr->ListDir(innerDirPath);
    return loadTableDir;
}

void Checkpoint::LoadDataset(const vector<string>& embNames, const vector<CkptDataType>& saveDataTypes,
                             const unique_ptr<CkptDataHandler>& dataHandler, CkptData& ckptData)
{
    for (const auto& embName : embNames) {
        auto dataDir{innerDirPath + dirSeparator + embName};
        for (const auto& saveDataType : saveDataTypes) {
            auto datasetPath{dataDir + dirSeparator + dataHandler->GetDataDirName(saveDataType)};

            auto datasetDir{datasetPath + dirSeparator + "slice" + dataFileType};
            auto attributeDir{datasetPath + dirSeparator + "slice" + attribFileType};

            CkptTransData transData;
            LOG_DEBUG("Start reading data from: {}", attributeDir);
            auto dataElmtBytes{dataHandler->GetDataElmtBytes(CkptDataType::ATTRIBUTE)};
            bool flag = saveDataType == CkptDataType::TABLE_2_THRESH || saveDataType == CkptDataType::HIST_REC;
            if (flag && !CheckFileExist(attributeDir)) {
                LOG_DEBUG("Path: {} does not exists, no need to read it. And set noFeatAdmitAndEvictData true",
                          attributeDir);
                ckptData.noFeatAdmitAndEvictData = true;
                continue;
            }
            ReadStream(transData, attributeDir, CkptDataType::ATTRIBUTE, dataElmtBytes);

            dataElmtBytes = dataHandler->GetDataElmtBytes(saveDataType);
            if (saveDataType == CkptDataType::EMB_DATA) {
                ReadStreamForEmbData(transData, datasetDir, dataElmtBytes, ckptData, embName);
                continue;
            } else {
                LOG_DEBUG("Start reading data from: {}", datasetDir);
                ReadStream(transData, datasetDir, saveDataType, dataElmtBytes);
            }

            LOG_DEBUG("Start loading data from: {} to data handler.", attributeDir);
            if ((saveDataType == CkptDataType::EMB_INFO)) {
                dataHandler->SetDatasetForLoadEmb(saveDataType, embName, transData, ckptData);
            } else {
                dataHandler->SetDataset(saveDataType, embName, transData);
            }
        }
    }
}

void Checkpoint::ReadStream(CkptTransData& transData, const string& dataDir, CkptDataType dataType,
                            uint32_t dataElmtBytes)
{
    if (dataElmtBytes == 0) {
        LOG_WARN("dataElmtBytes is 0, skip ReadStream");
        return;
    }

    CheckFileSystemPtr();

    size_t datasetSize = fileSystemPtr->GetFileSize(dataDir);
    auto resizeSize{datasetSize / dataElmtBytes};
    SetTransDataSize(transData, resizeSize, dataType);

    if (datasetSize % dataElmtBytes > 0) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::INVALID_ARGUMENT,
                           StringFormat("Data is missing or incomplete in load file: %s.", dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    ssize_t readBytesNum;
    if (int32TransSet.find(dataType) != int32TransSet.end()) {
        readBytesNum = fileSystemPtr->Read(dataDir, reinterpret_cast<char*>(transData.int32Arr.data()), datasetSize);
    } else if (int64TransSet.find(dataType) != int64TransSet.end()) {
        readBytesNum = fileSystemPtr->Read(dataDir, reinterpret_cast<char*>(transData.int64Arr.data()), datasetSize);
    } else if (dataType == CkptDataType::ATTRIBUTE) {
        readBytesNum = fileSystemPtr->Read(dataDir, reinterpret_cast<char*>(transData.attribute.data()), datasetSize);
    } else {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::LOGIC_ERROR, "Unknown CkptDataType.");
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    if (readBytesNum == -1) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::IO_ERROR,
                           StringFormat("Error: Load data failed. data type: %s. "
                                        "An error occurred while reading file: %s.",
                                        CkptDataTypeName(dataType).c_str(), dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
    if (readBytesNum != datasetSize) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::IO_ERROR,
                           StringFormat("Error: Load data failed. data type: %s. "
                                        "Expected to read %d bytes, but actually read %d bytes to file %s.",
                                        CkptDataTypeName(dataType).c_str(), datasetSize, readBytesNum,
                                        dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
}

void Checkpoint::ReadStreamForEmbData(CkptTransData& transData, const string& dataDir, uint32_t dataElmtBytes,
                                      CkptData& ckptData, string embName) const
{
    if (dataElmtBytes == 0) {
        LOG_WARN("dataElmtBytes is 0, skip ReadStreamForEmbData");
        return;
    }

    CheckFileSystemPtr();

    auto embDataOuterSize = transData.attribute.at(attribEmbDataOuterIdx);
    if (embDataOuterSize <= 0 || embDataOuterSize > MAX_VOCABULARY_SIZE) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::INVALID_ARGUMENT,
                           Logger::Format("Invalid embDataOuterSize:{}. The embDataOuterSize does not meet"
                               " the specified range ({}, {}).", embDataOuterSize, 0, MAX_VOCABULARY_SIZE));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    size_t datasetSize = fileSystemPtr->GetFileSize(dataDir);
    if (datasetSize % embDataOuterSize > 0 || datasetSize % dataElmtBytes > 0) {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::INVALID_ARGUMENT,
                           StringFormat("Data is missing or incomplete in load file: %s.", dataDir.c_str()));
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }

    auto loadHostEmbs = ckptData.hostEmbs;
    auto& dst = (*loadHostEmbs)[embName].embData;
    dst.reserve(embDataOuterSize);
}

void Checkpoint::SetTransDataSize(CkptTransData& transData, size_t datasetSize, CkptDataType dataType)
{
    if (int32TransSet.find(dataType) != int32TransSet.end()) {
        transData.int32Arr.resize(datasetSize);
    } else if (int64TransSet.find(dataType) != int64TransSet.end()) {
        transData.int64Arr.resize(datasetSize);
    } else if (dataType == CkptDataType::ATTRIBUTE) {
        transData.attribute.resize(datasetSize);
    } else {
        auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::LOGIC_ERROR, "Unknown CkptDataType.");
        LOG_ERROR(error.ToString());
        throw std::runtime_error(error.ToString());
    }
}

void Checkpoint::CheckFileSystemPtr() const
{
    if (fileSystemPtr != nullptr) {
        return;
    }
    auto error = Error(ModuleName::M_CHECK_POINT, ErrorType::NULL_PTR,
                       "FileSystemPtr is null pointer, please init file system pointer before using.");
    LOG_ERROR(error.ToString());
    throw std::runtime_error(error.ToString());
}
